Back to Blog | LearnableLoopAI.com | Portfolio of Projects |

versioninfo() ## Julia version

Julia Version 1.10.5
Commit 6f3fdf7b362 (2024-08-27 14:19 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 12 × Intel(R) Core(TM) i7-8700B CPU @ 3.20GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, skylake)
Threads: 1 default, 0 interactive, 1 GC (on 12 virtual cores)

import Pkg
Pkg.add(Pkg.PackageSpec(;name="RxInfer"))
Pkg.add("Plots")
Pkg.add("StableRNGs")
Pkg.add("LaTeXStrings")
Pkg.add("Distributions")

using RxInfer, Random, Plots, StableRNGs, LaTeXStrings, Distributions

   Resolving package versions...
  No Changes to `/workspaces/Kalman filtering and smoothing/Project.toml`
  No Changes to `/workspaces/Kalman filtering and smoothing/Manifest.toml`
   Resolving package versions...
  No Changes to `/workspaces/Kalman filtering and smoothing/Project.toml`
  No Changes to `/workspaces/Kalman filtering and smoothing/Manifest.toml`
   Resolving package versions...
  No Changes to `/workspaces/Kalman filtering and smoothing/Project.toml`
  No Changes to `/workspaces/Kalman filtering and smoothing/Manifest.toml`
   Resolving package versions...
  No Changes to `/workspaces/Kalman filtering and smoothing/Project.toml`
  No Changes to `/workspaces/Kalman filtering and smoothing/Manifest.toml`
   Resolving package versions...
  No Changes to `/workspaces/Kalman filtering and smoothing/Project.toml`
  No Changes to `/workspaces/Kalman filtering and smoothing/Manifest.toml`

Pkg.status()

Status `/workspaces/Kalman filtering and smoothing/Project.toml`
  [6e4b80f9] BenchmarkTools v1.8.0
  [31c24e10] Distributions v0.25.131
  [b964fa9f] LaTeXStrings v1.4.1
  [91a5bcdd] Plots v1.41.7
⌅ [86711068] RxInfer v3.10.1
  [860ef19b] StableRNGs v1.0.4
Info Packages marked with ⌅ have new versions available but compatibility constraints restrict them from upgrading. To see why use `status --outdated`

Kalman filtering and smoothing (Part 2)

This is an analysis of the RxInfer example at https://examples.rxinfer.com/categories/basic_examples/kalman_filtering_and_smoothing/
Some symbols have been changed
Some content has been added/modified
The preference is to make the math and code names align as much as possible
Spatial structure identifiers (e.g. vectors, matrices, cuboids)
- math
  - have a boldface e.g. \(\mathbf{x}\) or \(\boldsymbol{x}\) or \(\mathbf{A}\) or \(\mathbf{\mathbb{X}}\)
- code
  - have a underscore prefix e.g. _x or _A or _𝕏
- small case / cap case / blackboard case are generally used to discriminate between vectors, matrices, and cuboids
- start of alphabet / end of alphabet are generally used to discriminate between ‘system’ and ‘signal’
Time structure identifiers (i.e. time sequence/series)
- math
  - have a colon subscript e.g. \(x_:\) or \(x_{0:n}\) or \(x_{:n}\)
- code
  - have a ː (length mark) suffix e.g. xː
External (i.e. true, environment) states and parameters identifiers
- math
  - have a superscript * e.g. \(v^*\)
- code
  - have a superscript x e.g. vˣ
  - the ‘x’ in the code superscript is used to imitate the * superscript in the math

In the following set of examples the goal is to estimate hidden states of a Dynamical process where all hidden states are Gaussians.

We start our journey with a simple

1. multivariate Linear Gaussian State Space Model (LGSSM), which can be solved analytically. We then solve an
1. identification problem which does not have an analytical solution. Utimately, we show how RxInfer.jl can
1. deal with missing observations.

2 System Identification Problem

In this example we are going to attempt to run Bayesian inference and decouple two random-walk signals, which were combined into a single through some deterministic function f. We do not have access to the real values of these signals, but only to their combination. First, we create the sim_batch_data function that accepts f as an argument.

The function returns the real signals \(x^*_1\) and \(x^*_2\) for later comparison (we are not going to use them during inference) and their combined version \(y\) (we are going to use it as our observations during the inference). We also assume that \(y\) is corrupted with some measurement noise.

Combination 1: \(y = x^_1 + x^_2\)

In our first example, we are going to use a simple addition (+) as the function f. In general, it is impossible to decouple the signals \(x^*_1\) and \(x^*_2\) without strong priors, but we can try and see how good an inference can be. The + operation on two random variables also has a special meaning in the probabilistic inference, namely the convolution of pdf’s of the two random variables, and RxInfer treats it specially with many precomputed analytical rules, which may make the inference task easier. First, let us create a test dataset:

xˣ₁₍ₜ₋₁₎ = -20.0 ## x star 1 at t-1 (not done here as a component of a vector xˣ)
xˣ₂₍ₜ₋₁₎ = 20.0 ## x star 2 at t-1 (not done here as a component of a vector xˣ)
λˣ_x₁ = 0.1
λˣ_x₂ = 1.0
σˣ²ᵥ = 20.0 ## observation noise
T = 250

The Generative Process

State transition function (\(f_E\))

The state transition function provides the deterministic part of the state flow. The probabilistic part is provided by the system noise:

\[\dot{x}^*_{1t} = f_E(x^*_{1,t-1}) + \omega^*_{x1,t}\]

\[\dot{x}^*_{2t} = f_E(x^*_{2,t-1}) + \omega^*_{x2,t}\]

The *s indicate that the parameters and variables are true and not observed.

## state transition function
function fE(; xˣₜ₋₁)
    return xˣₜ₋₁
end
fE(xˣₜ₋₁=xˣ₁₍ₜ₋₁₎), fE(xˣₜ₋₁=xˣ₂₍ₜ₋₁₎)

(-20.0, 20.0)

Observation generation function (\(g_E\))

The observation generation function provides the deterministic part of the observation. The probabilistic part is provided by the observation noise:

\[y_{t} = g_E(x^*_{1t}, x^*_{2t}) + \sigma^{*2}_V = f(x^*_{1t}, x^*_{2t}) + \sigma^{*2}_V = x^*_{1t} + x^*_{2t} + \sigma^{*2}_V\]

The *s indicate that the parameters and variables are hidden and not observed.

## observation generation function
function gE(; f, xˣ₁ₜ, xˣ₂ₜ)
    return f(xˣ₁ₜ, xˣ₂ₜ)
end
gE(; f=+, xˣ₁ₜ=7.0, xˣ₂ₜ=8.0)

15.0

## Data comes from either a simulation/lab (sim|lab) OR from the field (fld)
## Data are handled either in batches (batch) OR online as individual points (point)
## Batch data accumulates either
    ## along the depth/examples dimension/axis (into the screen/page), OR
        ## typical for supervised & unsupervised learning
    ## along the time dimension/axis (down the screen page)
        ## typical for sequential decision learning (reinforcement learning & active inference)
function sim_batch_data(f, T; seed=123, xˣ₁₍ₜ₋₁₎, xˣ₂₍ₜ₋₁₎, λˣ_x₁, λˣ_x₂, σˣ²ᵥ)
    rng = StableRNG(seed)
    fE₁ː = Vector{Float64}(undef, T)
    xˣ₁ː = Vector{Float64}(undef, T)
    fE₂ː = Vector{Float64}(undef, T)
    xˣ₂ː = Vector{Float64}(undef, T)
    gEː = Vector{Float64}(undef, T)
    yː = Vector{Float64}(undef, T)
    for t in 1:T
        fE₁ː[t] = fE(xˣₜ₋₁=xˣ₁₍ₜ₋₁₎)
        xˣ₁ː[t] = rand(rng, Normal(fE₁ː[t], sqrt(1.0/λˣ_x₁)))
        fE₂ː[t] = fE(xˣₜ₋₁=xˣ₂₍ₜ₋₁₎)
        xˣ₂ː[t] = rand(rng, Normal(fE₂ː[t], sqrt(1.0/λˣ_x₂)))
        gEː[t] = gE(f=f, xˣ₁ₜ=xˣ₁ː[t], xˣ₂ₜ=xˣ₂ː[t])
        yː[t] = rand(rng, Normal(gEː[t], sqrt(σˣ²ᵥ)))
        xˣ₁₍ₜ₋₁₎ = xˣ₁ː[t]
        xˣ₂₍ₜ₋₁₎ = xˣ₂ː[t]
    end
    return xˣ₁ː, xˣ₂ː, yː
end

sim_batch_data (generic function with 1 method)

xˣ₁ː, xˣ₂ː, yː = sim_batch_data(
    +, T; 
    xˣ₁₍ₜ₋₁₎=xˣ₁₍ₜ₋₁₎, xˣ₂₍ₜ₋₁₎=xˣ₂₍ₜ₋₁₎, 
    λˣ_x₁=λˣ_x₁, λˣ_x₂=λˣ_x₂, 
    σˣ²ᵥ=σˣ²ᵥ);

xˣ₁ː

250-element Vector{Float64}:
 -20.401095992460988
 -20.669875759133884
 -22.790697414814776
 -27.104638760134062
 -31.248220640067366
 -26.079608998684243
 -27.11512896946521
 -30.294795287103234
 -32.698280013763174
 -34.850070378118104
 -36.54097463768623
 -37.561420452198526
 -30.47681847069751
   ⋮
 -58.883836692177226
 -60.76366280720403
 -55.502809787529735
 -54.42533877570402
 -52.38587391024251
 -52.85091503729493
 -56.78914098796256
 -55.32020135002518
 -57.28044133233906
 -59.6499029990486
 -59.74600957488878
 -59.32390477663594

xˣ₂ː

250-element Vector{Float64}:
 20.666885172487124
 20.57261116763893
 19.857870775719412
 18.858298515749894
 19.81453408886256
 20.425064540658344
 19.918312704314683
 19.372737596030216
 18.846506733028683
 19.028263308009453
 17.695641836921535
 19.601876559114125
 20.234315687020793
  ⋮
 51.44872861496441
 51.627760698406725
 51.67939447851273
 52.195724854526475
 52.15806523930519
 50.846307575145
 51.655335116038906
 50.10235490650183
 50.4547248217021
 51.8221505241864
 49.89530546131343
 49.41021524826118

yː

250-element Vector{Float64}:
  -5.353952716032312
   4.275970805667579
  -7.258265483265913
 -14.918691611380964
  -2.289681163005755
  -3.2075768172574284
  -2.526501156230419
  -7.968874122573138
 -15.558887021638972
 -18.922064684720848
 -19.917843395248976
 -12.92081654780483
  -6.814909043088918
   ⋮
  -1.5291300036990072
 -13.70504610995755
  -6.811199301966651
  -5.730263980041214
  -7.608774757258874
   0.8729760214719698
  -8.777046520647453
  -6.9699809935290045
  -8.267463283992132
  -1.8915218553802293
 -11.710160406427732
 -17.391946039775874

pl = plot(title="Underlying signals")
pl = plot!(pl, xˣ₁ː, label=L"x^*_1", c="red")
pl = plot!(pl, xˣ₂ː, label=L"x^*_2", c="blue")
pr = plot(title="Combined " * L"y = x^*_1 + x^*_2")
pr = scatter!(pr, yː, ms=3, color=:purple, label=L"y")
plot(pl, pr, size=(800, 300))

The Generative Model

To run inference, we need to create a probabilistic model: our beliefs about how our data could have been generated. For this we can use the @model macro from RxInfer.jl:

@model function identification_problem(f, yː, m₁₀, λ₁₀, a₁, b₁, m₂₀, λ₂₀, a₂, b₂, a_y, b_y)
    x₁₀ ~ Normal(mean=m₁₀, precision=λ₁₀)
    λ_x₁ ~ Gamma(shape=a₁, rate=b₁)
    x₂₀ ~ Normal(mean=m₂₀, precision=λ₂₀)
    λ_x₂ ~ Gamma(shape=a₂, rate=b₂)
    λ_y ~ Gamma(shape=a_y, rate=b_y)
    x₁₍ₜ₋₁₎ = x₁₀
    x₂₍ₜ₋₁₎ = x₂₀
    local x₁
    local x₂
    local z
    for t in 1:length(yː)
        x₁ː[t] ~ Normal(mean=x₁₍ₜ₋₁₎, precision=λ_x₁)
        x₂ː[t] ~ Normal(mean=x₂₍ₜ₋₁₎, precision=λ_x₂)
        zː[t] ~ f(x₁ː[t], x₂ː[t])
        yː[t] ~ Normal(mean=zː[t], precision=λ_y)
        x₁₍ₜ₋₁₎ = x₁ː[t]
        x₂₍ₜ₋₁₎ = x₂ː[t]
    end
end

RxInfer runs Bayesian inference as a variational optimisation procedure between the real solution and its variational proxy q. In our model specification we assumed noise components to be unknown, thus, we need to enforce a structured mean-field assumption for the variational family of distributions q. This inevitably reduces the accuracy of the result, but makes the task easier and allows for fast and analytical message passing-based variational inference:

constraints = @constraints begin
    q(x₁₀, x₂₀, x₁ː, x₂ː, λ_x₁, λ_x₂, λ_y, zː) = q(x₁ː, x₁₀, x₂ː, x₂₀, zː)q(λ_x₂)q(λ_x₁)q(λ_y)
end

Constraints: 
  q(x₁₀, x₂₀, x₁ː, x₂ː, λ_x₁, λ_x₂, λ_y, zː) = q(x₁ː, x₁₀, x₂ː, x₂₀, zː)q(λ_x₂)q(λ_x₁)q(λ_y)

The next step is to assign priors, initialise needed messages and marginals and call the inference function:

m₁₀, λ₁₀ = -20.0, 1.0
m₂₀, λ₂₀ = 20.0, 1.0

## We set relatively strong priors for random walk noise components
## and sort of vague prior for the noise of the observations
a₁, b₁ = 0.01, 0.01var(xˣ₁ː)
a₂, b₂ = 0.01, 0.01var(xˣ₂ː)
a_y, b_y = 1.0, 1.0

## We set relatively strong priors for messages
x₁_init = map(r -> NormalMeanPrecision(r, λ₁₀), reverse(range(-60, -20, length=T)))
x₂_init = map(r -> NormalMeanPrecision(r, λ₂₀), range(20, 60, length=T))

init = @initialization begin
    μ(x₁ː) = x₁_init
    μ(x₂ː) = x₂_init
    q(λ_x₁) = GammaShapeRate(a₁, b₁)
    q(λ_x₂) = GammaShapeRate(a₂, b₂)
    q(λ_y) = GammaShapeRate(a_y, b_y)
end

result = infer(
    model = identification_problem(
        f=+, 
        m₁₀=m₁₀, λ₁₀=λ₁₀, a₁=a₁, b₁=b₁, 
        m₂₀=m₂₀, λ₂₀=λ₂₀, a₂=a₂, b₂=b₂, 
        a_y=a_y, b_y=b_y),
    data  = (yː = yː,), 
    options = (limit_stack_depth = 500, ), 
    constraints = constraints, 
    initialization = init,
    iterations = 50
)

Inference results:
  Posteriors       | available for (zː, x₂₀, x₁₀, λ_x₂, λ_y, λ_x₁, x₂ː, x₁ː)

result.posteriors

Dict{Symbol, Vector} with 8 entries:
  :zː   => Vector{NormalWeightedMeanPrecision{Float64}}[[NormalWeightedMeanPrec…
  :x₂₀  => NormalWeightedMeanPrecision{Float64}[NormalWeightedMeanPrecision{Flo…
  :x₁₀  => NormalWeightedMeanPrecision{Float64}[NormalWeightedMeanPrecision{Flo…
  :λ_x₂ => GammaShapeRate{Float64}[GammaShapeRate{Float64}(a=125.01, b=128.149)…
  :λ_y  => GammaShapeRate{Float64}[GammaShapeRate{Float64}(a=126.0, b=261.766),…
  :λ_x₁ => GammaShapeRate{Float64}[GammaShapeRate{Float64}(a=125.01, b=132.229)…
  :x₂ː  => Vector{NormalWeightedMeanPrecision{Float64}}[[NormalWeightedMeanPrec…
  :x₁ː  => Vector{NormalWeightedMeanPrecision{Float64}}[[NormalWeightedMeanPrec…

Let’s examine our inference results:

λ_x₁_marginals = result.posteriors[:λ_x₁]
λ_x₂_marginals = result.posteriors[:λ_x₂]
λ_y_marginals = result.posteriors[:λ_y]

z_marginals = result.posteriors[:zː]
x₁_marginals = result.posteriors[:x₁ː]
x₂_marginals = result.posteriors[:x₂ː];

p1 = plot(legend=:bottomleft, title="Estimated hidden signals")
p2 = plot(legend=:bottomright, title="Estimated combined signals")

p1 = plot!(p1, xˣ₁ː, label="Real hidden xˣ₁", c="red", linestyle=:dash)
p1 = plot!(p1, mean.(x₁_marginals[end]), ribbon=var.(x₁_marginals[end]), label="Estimated x₁", color=:pink)

p1 = plot!(p1, xˣ₂ː, label="Real hidden xˣ₂", color=:blue, linestyle=:dash)
p1 = plot!(p1, mean.(x₂_marginals[end]), ribbon=var.(x₂_marginals[end]), label="Estimated x₂", color=:lightblue)

p2 = scatter!(p2, yː, label="Observations", ms=2, alpha=0.5, color=:orange)
p2 = plot!(p2, mean.(z_marginals[end]), ribbon=std.(z_marginals[end]), label="Combined estimated signal", color="plum")

plot(p1, p2, size=(800, 300))

The inference results are not so bad, even though RxInfer missed the correct values of the signals between 100 and 150.

Combination 2: \(y = \mathrm{min}(x^_1, x^_2)\)

In this example we use a slightly more complex function, for which RxInfer does not have precomputed analytical message update rules. We are going to attempt to run Bayesian inference with min as a combination function. Note, however, that directly using min may cause problems for the built-in approximation methods as it has zero partial derviates with respect to all but one of the variables.

The Generative Process

xˣ₁₍ₜ₋₁₎ = 0.0 ## x star 1 at t-1 (not done here as a component of a vector xˣ)
xˣ₂₍ₜ₋₁₎ = 0.0 ## x star 2 at t-1 (not done here as a component of a vector xˣ)
λˣ_x₁ = 1.0
λˣ_x₂ = 1.0
σˣ²ᵥ = 1.0 ## observation noise
T = 200

min_xˣ₁ː, min_xˣ₂ː, min_yː = sim_batch_data(
    min, T; 
    seed=1, 
    # seed=123, 
    xˣ₁₍ₜ₋₁₎=xˣ₁₍ₜ₋₁₎, xˣ₂₍ₜ₋₁₎=xˣ₂₍ₜ₋₁₎, 
    λˣ_x₁=λˣ_x₁, λˣ_x₂=λˣ_x₂, 
    σˣ²ᵥ=σˣ²ᵥ);

min_xˣ₁ː

200-element Vector{Float64}:
 -0.5325200748641231
 -1.3760574787383777
 -0.34376400643438854
 -0.09683736607523746
 -1.9500360726184451
 -1.7178925137280041
 -0.5114889405789891
  0.9799054887699994
  1.4738614452691385
  0.21583988918877006
  0.5371882549212899
  0.5915765895874429
  1.0901880014255332
  ⋮
 11.872449055238848
 11.226459848650638
 11.84602940087337
 11.635994703121376
 10.924903891533114
 12.101086505558115
 12.23488278202865
 14.24470601685975
 16.68915952947985
 15.21355198384941
 17.49518040732118
 16.88853455216658

min_xˣ₂ː

200-element Vector{Float64}:
  0.098465514284785
 -1.9153638413463216
 -1.2397977422271311
 -2.740326889546081
 -3.0775177160930336
 -2.5211723474749315
 -3.1417645469162077
 -3.8175642158895347
 -3.403893528229725
 -3.944508580769285
 -3.2562253933941165
 -2.5279231067335415
 -3.3863346347473797
  ⋮
  2.7276444782559133
  2.2683338897721086
  2.690780210483017
  2.9447026640735814
  3.343236550774581
  1.7153229626545394
  3.1612860855506497
  3.464740837552133
  4.966929654823941
  5.590221124459778
  6.156878620294983
  5.219136120500052

min_yː

200-element Vector{Float64}:
  0.22036644726040033
 -2.1851714825434487
 -1.1576369146831669
 -2.411694880333956
 -3.0418551997138934
 -3.5489933567433987
 -2.529456341036421
 -4.568787926083518
 -2.8472451081518817
 -4.672624206001446
 -3.646165280810205
 -1.1964331836419906
 -2.3952985765408528
  ⋮
  2.137866865028408
  3.1158574599027653
  3.1908954000708625
  2.7076738085196226
  3.5105327668426
  2.1476892369392173
  4.2435567676856065
  4.165753448418778
  3.584702727331395
  6.14602156701984
  7.5620630212942235
  6.086359625031924

pl = plot(title = "Underlying signals")
pl = plot!(pl, min_xˣ₁ː, label=L"x^*_1", c=:red)

pl = plot!(pl, min_xˣ₂ː, label=L"x^*_2", c=:blue)

pr = plot(title = "Combined " * L"y = \mathrm{min}(x^*_1, x^*_2)")
pr = scatter!(pr, min_yː, ms=3, color=:purple, label=L"y")

plot(pl, pr, size = (800, 300))

The Generative Model

We generate data with the min function directly, however we model it with a somewhat smoothed version:

## Smoothed version of `min` without zero-ed derivatives
function smooth_min(x, y)    
    if x < y
        return x + 1e-4*y
    else
        return y + 1e-4*x
    end
end

smooth_min (generic function with 1 method)

RxInfer supports arbitrary nonlinear functions, but it requires an explicit approximation method specification. That can be achieved with the built-in @meta macro:

min_meta = @meta begin
    ## In this example we are going to use a simple `Linearization` method
    smooth_min() -> Linearization()
end

Meta: 
  smooth_min() -> Linearization()

min_m₁₀, min_λ₁₀ = -1.0, 1.0
min_m₂₀, min_λ₂₀ = 1.0, 1.0

min_a₁, min_b₁ = 1.0, 1.0
min_a₂, min_b₂ = 1.0, 1.0
min_a_y, min_b_y = 1.0, 1.0

init = @initialization begin
    μ(x₁ː) = NormalMeanPrecision(min_m₁₀, min_λ₁₀)
    μ(x₂ː) = NormalMeanPrecision(min_m₂₀, min_λ₂₀)
    q(λ_x₁) = GammaShapeRate(min_a₁, min_b₁)
    q(λ_x₂) = GammaShapeRate(min_a₂, min_b₂)
    q(λ_y) = GammaShapeRate(min_a_y, min_b_y)
end

min_result = infer(
    model = identification_problem(
        f=smooth_min, 
        m₁₀=min_m₁₀, λ₁₀=min_λ₁₀, a₁=min_a₁, b₁=min_b₁, 
        m₂₀=min_m₂₀, λ₂₀=min_λ₂₀, a₂=min_a₂, b₂=min_b₂, 
        a_y=min_a_y, b_y=min_b_y),
    data  = (yː = min_yː,), 
    options = (limit_stack_depth = 500, ), 
    constraints = constraints, 
    initialization = init,
    meta = min_meta,
    iterations = 50
)

Inference results:
  Posteriors       | available for (zː, x₂₀, x₁₀, λ_x₂, λ_y, λ_x₁, x₂ː, x₁ː)

min_result.posteriors

Dict{Symbol, Vector} with 8 entries:
  :zː   => Vector{NormalWeightedMeanPrecision{Float64}}[[NormalWeightedMeanPrec…
  :x₂₀  => NormalWeightedMeanPrecision{Float64}[NormalWeightedMeanPrecision{Flo…
  :x₁₀  => NormalWeightedMeanPrecision{Float64}[NormalWeightedMeanPrecision{Flo…
  :λ_x₂ => GammaShapeRate{Float64}[GammaShapeRate{Float64}(a=101.0, b=51.2907),…
  :λ_y  => GammaShapeRate{Float64}[GammaShapeRate{Float64}(a=101.0, b=200.148),…
  :λ_x₁ => GammaShapeRate{Float64}[GammaShapeRate{Float64}(a=101.0, b=51.2789),…
  :x₂ː  => Vector{NormalWeightedMeanPrecision{Float64}}[[NormalWeightedMeanPrec…
  :x₁ː  => Vector{NormalWeightedMeanPrecision{Float64}}[[NormalWeightedMeanPrec…

min_λ_x₁_marginals = min_result.posteriors[:λ_x₁]
min_λ_x₂_marginals = min_result.posteriors[:λ_x₂]
min_λ_y_marginals = min_result.posteriors[:λ_y]

min_z_marginals = min_result.posteriors[:zː]
min_x₁_marginals = min_result.posteriors[:x₁ː]
min_x₂_marginals = min_result.posteriors[:x₂ː];

p1 = plot(legend=:bottomleft, title="Estimated hidden signals")
p2 = plot(legend=:bottomright, title="Estimated combined signals")

p1 = plot!(p1, min_xˣ₁ː, label="Real hidden xˣ₁", c="red", linestyle=:dash)
p1 = plot!(p1, mean.(min_x₁_marginals[end]), ribbon=var.(min_x₁_marginals[end]), label="Estimated x₁", color=:pink)

p1 = plot!(p1, min_xˣ₂ː, label="Real hidden xˣ₂", color=:blue, linestyle=:dash)
p1 = plot!(p1, mean.(min_x₂_marginals[end]), ribbon=var.(min_x₂_marginals[end]), label="Estimated x₂", color=:lightblue)

p2 = scatter!(p2, min_yː, label="Observations", ms=2, alpha=0.5, color=:orange)
p2 = plot!(p2, mean.(min_z_marginals[end]), ribbon=std.(min_z_marginals[end]), label="Combined estimated signal", color="plum")

plot(p1, p2, size=(800, 300))

As we can see inference with the min function is significantly harder. Even though the combined signal has been inferred with high precision the underlying v and w signals are barely inferred. This may be expected, since the min function essentially destroy the information about one of the signals, thus, making it impossible to decouple two seemingly identical random walk signals. The only one inferred signal is the one which is lower and we have no inference information about the signal which is above. It might be possible to infer the states, however, with more informative priors and structural information about two different signals (e.g. if these are not random walks).

Online (filtering) identification: \(y = \mathrm{min}(x^_1, x^_2)\)

Another way to approach this problem is to use online (filtering) inference procedure from RxInfer, but for that we also need to modify our model specification a bit:

The Generative Process

Next step is to generate our dataset and to run the actual inference procedure! For that we use the infer function with autoupdates keyword:

xˣ₁₍ₜ₋₁₎ = 1.0 ## x star 1 at t-1 (not done here as a component of a vector xˣ)
xˣ₂₍ₜ₋₁₎ = -1.0 ## x star 2 at t-1 (not done here as a component of a vector xˣ)
λˣ_x₁ = 1.0
λˣ_x₂ = 1.0
σˣ²ᵥ = 1.0 ## observation noise
T = 300

rx_xˣ₁ː, rx_xˣ₂ː, rx_yː = sim_batch_data(
    min, T; 
    seed=1, 
    xˣ₁₍ₜ₋₁₎=xˣ₁₍ₜ₋₁₎, xˣ₂₍ₜ₋₁₎=xˣ₂₍ₜ₋₁₎, 
    λˣ_x₁=λˣ_x₁, λˣ_x₂=λˣ_x₂, 
    σˣ²ᵥ=σˣ²ᵥ);

rx_xˣ₁ː

300-element Vector{Float64}:
  0.4674799251358769
 -0.3760574787383777
  0.6562359935656115
  0.9031626339247625
 -0.9500360726184451
 -0.7178925137280042
  0.4885110594210108
  1.9799054887699992
  2.4738614452691383
  1.2158398891887698
  1.5371882549212896
  1.5915765895874425
  2.090188001425533
  ⋮
 13.384009842250418
 13.053540348684535
 13.595007553814803
 14.715705062022895
 15.13350407386025
 14.966386749228318
 15.005789040066146
 14.522153754344416
 13.542981339795244
 14.554317190702813
 12.879159788609714
 12.533966078243003

rx_xˣ₂ː

300-element Vector{Float64}:
  -0.901534485715215
  -2.9153638413463216
  -2.239797742227131
  -3.740326889546081
  -4.077517716093034
  -3.521172347474932
  -4.1417645469162085
  -4.8175642158895355
  -4.403893528229726
  -4.944508580769286
  -4.256225393394117
  -3.527923106733542
  -4.38633463474738
   ⋮
  -1.3708151099757697
  -2.0614994788855805
  -2.7625770269891565
  -3.127596543119536
  -2.4904468534292046
  -2.41950024394849
  -2.8024522739003057
  -5.079038511589717
  -6.533981992409512
  -9.158210504701648
 -10.071390396703338
 -10.363573258463711

rx_yː

300-element Vector{Float64}:
  -0.14864796359069155
  -3.1851714825434487
  -2.1576369146831667
  -3.411694880333956
  -4.041855199713893
  -4.5489933567433996
  -3.529456341036422
  -5.568787926083519
  -3.8472451081518826
  -5.672624206001447
  -4.646165280810205
  -2.1964331836419912
  -3.3952985765408528
   ⋮
  -2.096338290076389
  -0.8576465355912335
  -3.3898995696931755
  -0.8757311338953824
  -1.7386965707479658
  -1.7192256893794404
  -2.799924492306463
  -3.534637117106459
  -5.694563357428694
 -10.969410855554415
  -9.983476577762598
 -11.66386107476468

pl = plot(title="Underlying signals")
pl = plot!(pl, rx_xˣ₁ː, label=L"x^*_1", c=:red)

pl = plot!(pl, rx_xˣ₂ː, label=L"x^*_2", c=:blue)

pr = plot(title = "Combined " * L"y = \mathrm{min}(x^*_1, x^*_2)")
pr = scatter!(pr, rx_yː, ms=3, color=:purple, label=L"y")

plot(pl, pr, size=(800, 300))

The Generative Model

Another way to approach to this problem is to use online (filtering) inference procedure from RxInfer, but for that we also need to modify our model specification a bit:

@model function rx_identification(f, y, m₁₀, λ₁₀, a₁, b₁, m₂₀, λ₂₀, a₂, b₂, a_y, b_y)
    ## We are going to continuosly update our priors
    ## based on new posteriors
    x₁₀ ~ Normal(mean=m₁₀, precision=λ₁₀)
    λ_x₁ ~ Gamma(shape=a₁, rate=b₁)
    x₂₀ ~ Normal(mean=m₂₀, precision=λ₂₀)
    λ_x₂ ~ Gamma(shape=a₂, rate=b₂)
    λ_y ~ Gamma(shape=a_y, rate=b_y)
    
    x₁ ~ Normal(mean= x₁₀, precision= λ_x₁)
    x₂ ~ Normal(mean = x₂₀, precision = λ_x₂)

    z ~ f(x₁, x₂)
    y ~ Normal(mean=z, precision=λ_y)
end

We impose structured mean-field assumption for this model as well:

rx_constraints = @constraints begin
    q(x₁₀, x₁, x₂₀, x₂, λ_x₁, λ_x₂, λ_y, z) = q(x₁₀, x₁)q(x₂, x₂₀)q(λ_x₂)q(λ_x₁)q(z)q(λ_y)
end

Constraints: 
  q(x₁₀, x₁, x₂₀, x₂, λ_x₁, λ_x₂, λ_y, z) = q(x₁₀, x₁)q(x₂, x₂₀)q(λ_x₂)q(λ_x₁)q(z)q(λ_y)

Online inference in the RxInfer supports the @autoupdates specification, which tells inference procedure how to update priors based on new computed posteriors:

autoupdates = @autoupdates begin
    m₁₀, λ₁₀ = mean_precision(q(x₁))
    m₂₀, λ₂₀ = mean_precision(q(x₂))
    a₁ = shape(q(λ_x₁))
    b₁ = rate(q(λ_x₁))
    a_y = shape(q(λ_y))
    b_y = rate(q(λ_y))
    a₂ = shape(q(λ_x₂)) 
    b₂ = rate(q(λ_x₂))
end

@autoupdates begin
    (m₁₀, λ₁₀) = mean_precision(q(x₁))
    (m₂₀, λ₂₀) = mean_precision(q(x₂))
    a₁ = shape(q(λ_x₁))
    b₁ = rate(q(λ_x₁))
    a_y = shape(q(λ_y))
    b_y = rate(q(λ_y))
    a₂ = shape(q(λ_x₂))
    b₂ = rate(q(λ_x₂))
end

As previously we need to define the @meta structure that specifies the approximation method for the nonlinear function smooth_min (f in the model specification):

rx_meta = @meta begin 
    smooth_min() -> Linearization()
end

Meta: 
  smooth_min() -> Linearization()

init = @initialization begin
    q(x₁) = NormalMeanVariance(2.0, 1.0)
    q(x₂)= NormalMeanVariance(-2.0, 1.0)
    q(λ_x₁) = GammaShapeRate(1.0, 1.0) 
    q(λ_x₂) = GammaShapeRate(1.0, 1.0) 
    q(λ_y) = GammaShapeRate(1.0, 20.0)
end

engine = infer(
    model         = rx_identification(f=smooth_min),
    constraints   = rx_constraints,
    data          = (y = rx_yː,),
    autoupdates   = autoupdates,
    meta          = rx_meta,
    returnvars    = (:x₁, :x₂, :λ_x₁, :λ_x₂, :λ_y, :z),
    keephistory   = 1000,
    historyvars   =  KeepLast(),
    initialization = init,
    iterations    = 10,
    free_energy = true, 
    free_energy_diagnostics = nothing,
    autostart     = true,
)

RxInferenceEngine:
  Posteriors stream    | enabled for (λ_x₂, λ_y, λ_x₁, x₂, z, x₁)
  Free Energy stream   | enabled
  Posteriors history   | available for (x₂₀, x₁₀, λ_x₂, λ_y, λ_x₁, x₂, z, x₁)
  Free Energy history  | available
  Enabled events       | [  ]

engine.posteriors

Dict{Symbol, ProxyObservable{Any, Rocket.ScheduledSource{Marginal, PendingScheduler, ProxyObservable{Marginal, ReactiveMP.MarginalObservable, Rocket.FilterProxy{ReactiveMP.var"#31#32"}}}, Rocket.MapProxy{Marginal, RxInfer.var"#227#243"{DefaultPostprocess}}}} with 6 entries:
  :λ_x₂ => ProxyObservable(Any, MapProxy(Marginal))
  :λ_y  => ProxyObservable(Any, MapProxy(Marginal))
  :λ_x₁ => ProxyObservable(Any, MapProxy(Marginal))
  :x₂   => ProxyObservable(Any, MapProxy(Marginal))
  :z    => ProxyObservable(Any, MapProxy(Marginal))
  :x₁   => ProxyObservable(Any, MapProxy(Marginal))

rx_z_marginals = engine.history[:z]
rx_x₁_marginals = engine.history[:x₁]
rx_x₂_marginals = engine.history[:x₂];

p1 = plot(legend=:bottomleft, title="Estimated hidden signals")
p2 = plot(legend=:bottomright, title="Estimated combined signals")

p1 = plot!(p1, rx_xˣ₁ː, label="Real hidden xˣ₁", c=:red, linestyle=:dash)
p1 = plot!(p1, mean.(rx_x₁_marginals), ribbon=var.(rx_x₁_marginals), label="Estimated x₁", c=:pink)

p1 = plot!(p1, rx_xˣ₂ː, label="Real hidden xˣ₂", c=:blue, linestyle=:dash)
p1 = plot!(p1, mean.(rx_x₂_marginals), ribbon=var.(rx_x₂_marginals), label="Estimated x₂", c=:lightblue)

p2 = scatter!(p2, rx_yː, label="Observations", ms=2, alpha=0.5, color=:orange)
p2 = plot!(p2, mean.(rx_z_marginals), ribbon=std.(rx_z_marginals), label="Combined estimated signal", color=:plum)

plot(p1, p2, size=(800, 300))

The results are quite similar to the smoothing case and, as we can see, one of the random walk is again in the “disabled” state, does not infer anything and simply increases its variance (which is expected for the random walk).

Kalman filtering and smoothing (Part 2)

2 System Identification Problem

Combination 1: \(y = x^*_1 + x^*_2\)

The Generative Process

State transition function (\(f_E\))

Observation generation function (\(g_E\))

The Generative Model

Combination 2: \(y = \mathrm{min}(x^*_1, x^*_2)\)

The Generative Process

The Generative Model

Online (filtering) identification: \(y = \mathrm{min}(x^*_1, x^*_2)\)

The Generative Process

The Generative Model

Combination 1: \(y = x^_1 + x^_2\)

Combination 2: \(y = \mathrm{min}(x^_1, x^_2)\)

Online (filtering) identification: \(y = \mathrm{min}(x^_1, x^_2)\)